# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import collections

import os
import glob
import pandas

import logging
import functools
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = 16, 7
plt.rcParams['font.family'] = ['Dejavu Sans Mono']
plt.rcParams['font.sans-serif'] = ['Dejavu Sans Mono']
plt.rcParams['font.monospace'] = ['Dejavu Sans Mono']
plt.rcParams['font.size'] = 10
plt.rcParams['legend.fontsize'] = 10
plt.rcParams['xtick.labelsize'] = 10
plt.rcParams['ytick.labelsize'] = 10

from coin.base.datetime_util import to_timestamp_int

import coin.strategy.mm.tool.archive_base as abase


def main(argv):
  files = glob.glob("oe_stat_intra/*.csv")
  dfs = [pandas.read_csv(file) for file in files]
  dfdict = collections.defaultdict(list)

  for df in dfs:
    dfdict[df['currency'].iloc[-1]].append(df.iloc[1:, 1:])

  os.makedirs("pictest", exist_ok=True)
  ccys = ['BTC', 'ETH', 'XRP', 'LTC', 'BCHABC', 'EOS', 'BCHSV', 'ETC', 'TRX']
  vals = [8641.69, 172.32, 0.240011, 57.16, 231.54, 2.91, 88.8, 4.8, 0.013211]
  mtmdict = dict(zip(ccys, vals))

  for ccy in dfdict.keys():
    df_final = functools.reduce(
        lambda left, right: pandas.merge(left, right, on='ts_sec'), dfdict[ccy])[20:]
    print(ccy, df_final.head())
    pos = None
    bal = None
    for poscol in ['pos', 'pos_x', 'pos_y']:
      if poscol not in df_final:
        continue
      try:
        seri = df_final[poscol].sum(1)
      except Exception:
        seri = df_final[poscol]

      if pos is None:
        pos = seri
      else:
        pos += seri
    for balcol in ['balance', 'balance_x', 'balance_y']:
      if balcol not in df_final:
        continue

      try:
        seri = df_final[balcol].sum(1)
      except Exception:
        seri = df_final[balcol]

      if bal is None:
        bal = seri
      else:
        bal += seri

    bal *= 1.5

    mtmp = mtmdict[ccy]

    tidx = pandas.DatetimeIndex(df_final['ts_sec'] * 1e9)
    # plt.plot(tidx, pos, 'r-', lw=0.2)
    # plt.ylabel(f'{ccy} pos')
    # plt.twinx()
    plt.plot(tidx, bal, 'b-')
    plt.ylabel(f'{ccy} bal')
    loss = bal.iloc[0] - bal.min()
    plt.title(f'{loss : .2f} {ccy} * {mtmp} = {loss * mtmp : .2f} USD')
    plt.savefig(f'pictest/{loss*mtmp:.0f}_{ccy}.png')
    plt.close()


class Incident(object):
  def __init__(self, products, start_ts=None, end_ts=None):
    self._products = products
    self._bbos = {product: [] for product in products}
    self._instrs = {product: [] for product in products}
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))
      if hasattr(book_builder, "subscribe_instrument"):
        book_builder.subscribe_instrument(product, functools.partial(self.on_instrument, product))

  def on_instrument(self, product, instrument):
    if hasattr(instrument, 'open_interest'):
      open_interest = instrument.open_interest
      funding_rate = instrument.funding_rate
      mark_price = instrument.mark_price
      indicative_funding_rate = instrument.indicative_funding_rate
    else:
      return

    self._instrs[product].append({
        'timestamp': instrument.timestamp,
        'open_interest': open_interest,
        'funding_rate': funding_rate,
        'indicative_funding_rate': indicative_funding_rate,
        'mark_price': mark_price,
    })

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return

    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
    except (AttributeError, IndexError):
      return

    self._bbos[product].append({
        'timestamp': book.timestamp, 'bid0p': bid0_price, 'ask0p': ask0_price
    })

  def on_trade(self, product, trade):
    pass


def main2(argv):
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('symbol', None, 'symbol name')

  flags.FLAGS.exchange = 'Okex'
  flags.FLAGS.market_type = 'Futures'
  flags.FLAGS.trading_date = '20190924'
  flags.FLAGS.use_feed_cache = True
  flags.FLAGS.time_range = '18-24'

  ccys = ['BTC', 'ETH', 'XRP', 'LTC', 'BCHABC', 'EOS', 'BCHSV', 'ETC', 'TRX']

  for ccy in ccys:
    flags.FLAGS.symbol = f'{ccy}-USD.QUARTER'
    trading_dates = abase.get_trading_dates(flags.FLAGS.trading_date)
    fsflags = abase.get_feed_subscription_flags()
    baseflags = abase.get_base_flags()

    products = abase.symbols_to_products([flags.FLAGS.symbol], baseflags)
    start_time, end_time = abase.get_time_range(baseflags)

    strategy = Incident(products)

    abase.run_from_archive_base(baseflags=baseflags,
                                feed_sub_flags_proto=fsflags,
                                products=products,
                                on_book_reset=strategy.on_book_reset,
                                api_version_override={
                                    ('Futures', 'Okex'): 'v3',
                                    ('Spot', 'Okex'): 'v3',
                                })

    bbodf = pandas.DataFrame(strategy._bbos[products[0]])
    instrdf = pandas.DataFrame(strategy._instrs[products[0]])

    plt.plot(pandas.DatetimeIndex(bbodf['timestamp']),
             0.5 * (bbodf['ask0p'] + bbodf['bid0p']),
             'r-',
             pandas.DatetimeIndex(instrdf['timestamp']),
             instrdf['mark_price'],
             'b:',
             lw=0.5)
    plt.savefig(f"pictest/basis_{ccy}.png")
    plt.close()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
  # app.run(main2)
